{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_BoKmYo6oUci"
   },
   "source": [
    "License\n",
    "\n",
    "```\n",
    "Copyright (c) Facebook, Inc. and its affiliates.\n",
    "\n",
    "This source code is licensed under the MIT license found in the\n",
    "LICENSE file in the root directory of this source tree.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gsrdt9HooN9K"
   },
   "source": [
    "# Using CompilerGym environments with RLlib\n",
    "\n",
    "In this notebook we will use [RLlib](https://docs.ray.io/en/master/rllib.html) to train an agent for CompilerGym's [LLVM environment](https://facebookresearch.github.io/CompilerGym/llvm/index.html). RLlib is a popular library for scalable reinforcement learning, built on [Ray](https://docs.ray.io/en/master/index.html). It provides distributed implementations of several standard reinforcement learning algorithms.\n",
    "\n",
    "Our goal is not to produce the best agent, but to demonstrate how to integrate CompilerGym with RLlib. It will take about 20 minutes to work through. Let's get started!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W8CSxbx5ovuF"
   },
   "source": [
    "## Installation\n",
    "\n",
    "We'll begin by installing the `compiler_gym` and `ray` packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cT3QDtxbf3Cr",
    "outputId": "399fced8-ec0c-4745-cf29-74dac40429f4"
   },
   "outputs": [],
   "source": [
    "!pip install compiler_gym 'ray[default,rllib]' &>/dev/null || echo \"Install failed!\"\n",
    "\n",
    "# Print the versions of the libraries that we are using:\n",
    "import compiler_gym\n",
    "import ray\n",
    "\n",
    "print(\"compiler_gym version:\", compiler_gym.__version__)\n",
    "print(\"ray version:\", ray.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5Y2bt7GttpQ3"
   },
   "source": [
    "## Defining an Environment\n",
    "\n",
    "Next we will define the environment to use for our experiments. For the purposes of a simple demo we will apply two simplifying constraints to CompilerGym's LLVM environment:\n",
    "\n",
    "1. We will use only a small subset of the command line flag action space.\n",
    "2. We will clip the length of episodes to a maximum number of steps.\n",
    "\n",
    "To make things simple we will define a `make_env()` helper function to create our environment, and use the [compiler_gym.wrappers](https://facebookresearch.github.io/CompilerGym/compiler_gym/wrappers.html) API to implement these constraints. There is quite a lot going on in this cell, be sure to read through the comments for an explanation of what is going on!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bNeq2sArf5pa"
   },
   "outputs": [],
   "source": [
    "from compiler_gym.wrappers import ConstrainedCommandline, TimeLimit\n",
    "from ray import tune\n",
    "\n",
    "def make_env() -> compiler_gym.envs.CompilerEnv:\n",
    "    \"\"\"Make the reinforcement learning environment for this experiment.\"\"\"\n",
    "    # We will use LLVM as our base environment. Here we specify the observation\n",
    "    # space from this paper: https://arxiv.org/pdf/2003.00671.pdf and the total\n",
    "    # IR instruction count as our reward space, normalized against the \n",
    "    # performance of LLVM's -Oz policy.\n",
    "    env = compiler_gym.make(\n",
    "        \"llvm-v0\",\n",
    "        observation_space=\"Autophase\",\n",
    "        reward_space=\"IrInstructionCountOz\",\n",
    "    )\n",
    "    # Here we constrain the action space of the environment to use only a \n",
    "    # handful of command line flags from the full set. We do this to speed up\n",
    "    # learning by pruning the action space by hand. This also limits the \n",
    "    # potential improvements that the agent can achieve compared to using the \n",
    "    # full action space.\n",
    "    env = ConstrainedCommandline(env, flags=[\n",
    "        \"-break-crit-edges\",\n",
    "        \"-early-cse-memssa\",\n",
    "        \"-gvn-hoist\",\n",
    "        \"-gvn\",\n",
    "        \"-instcombine\",\n",
    "        \"-instsimplify\",\n",
    "        \"-jump-threading\",\n",
    "        \"-loop-reduce\",\n",
    "        \"-loop-rotate\",\n",
    "        \"-loop-versioning\",\n",
    "        \"-mem2reg\",\n",
    "        \"-newgvn\",\n",
    "        \"-reg2mem\",\n",
    "        \"-simplifycfg\",\n",
    "        \"-sroa\",\n",
    "    ])\n",
    "    # Finally, we impose a time limit on the environment so that every episode\n",
    "    # for 5 steps or fewer. This is because the environment's task is continuous\n",
    "    # and no action is guaranteed to result in a terminal state. Adding a time\n",
    "    # limit means we don't have to worry about learning when an agent should \n",
    "    # stop, though again this limits the potential improvements that the agent\n",
    "    # can achieve compared to using an unbounded maximum episode length.\n",
    "    env = TimeLimit(env, max_episode_steps=5)\n",
    "    return env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Eg84RXsDt5ey",
    "outputId": "8d1fa820-348e-400e-de6c-0a6d9a14d9bb"
   },
   "outputs": [],
   "source": [
    "# Let's create an environment and print a few attributes just to check that we \n",
    "# have everything set up the way that we would like.\n",
    "with make_env() as env:\n",
    "    print(\"Action space:\", env.action_space)\n",
    "    print(\"Observation space:\", env.observation_space)\n",
    "    print(\"Reward space:\", env.reward_space)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "39k80F8itTT_"
   },
   "source": [
    "## Datasets\n",
    "\n",
    "Now that we have an environment, we will need a set of programs to train on. In CompilerGym, these programs are called *benchmarks*. CompilerGym ships with [several sets of benchmarks](https://facebookresearch.github.io/CompilerGym/llvm/index.html#datasets). Here we will take a handful of benchmarks from the `npb-v0` dataset for training. We will then further divide this set into training and validation sets. We will use `chstone-v0` as a holdout test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IV6-adNIhoGS",
    "outputId": "23afc569-58d1-443a-a4a7-e72b703152cc"
   },
   "outputs": [],
   "source": [
    "from itertools import islice\n",
    "\n",
    "with make_env() as env:\n",
    "  # The two datasets we will be using:\n",
    "  npb = env.datasets[\"npb-v0\"]\n",
    "  chstone = env.datasets[\"chstone-v0\"]\n",
    "\n",
    "  # Each dataset has a `benchmarks()` method that returns an iterator over the\n",
    "  # benchmarks within the dataset. Here we will use iterator sliceing to grab a \n",
    "  # handful of benchmarks for training and validation.\n",
    "  train_benchmarks = list(islice(npb.benchmarks(), 55))\n",
    "  train_benchmarks, val_benchmarks = train_benchmarks[:50], train_benchmarks[50:]\n",
    "  # We will use the entire chstone-v0 dataset for testing.\n",
    "  test_benchmarks = list(chstone.benchmarks())\n",
    "\n",
    "print(\"Number of benchmarks for training:\", len(train_benchmarks))\n",
    "print(\"Number of benchmarks for validation:\", len(val_benchmarks))\n",
    "print(\"Number of benchmarks for testing:\", len(test_benchmarks))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZW0sfMhjv8Kg"
   },
   "source": [
    "## Registering the environment with RLlib\n",
    "\n",
    "Now that we have our environment and training benchmarks, we can register the environment for use with RLlib. To do this we will define a second `make_training_env()` helper that uses the [CycleOverBenchmarks](https://facebookresearch.github.io/CompilerGym/compiler_gym/wrappers.html#compiler_gym.wrappers.CycleOverBenchmarks) wrapper to ensure that the environment uses all of the training benchmarks. We then call `tune.register_env()`, assining the environment a name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-UgFKvTkv64L"
   },
   "outputs": [],
   "source": [
    "from compiler_gym.wrappers import CycleOverBenchmarks\n",
    "\n",
    "def make_training_env(*args) -> compiler_gym.envs.CompilerEnv:\n",
    "  \"\"\"Make a reinforcement learning environment that cycles over the\n",
    "  set of training benchmarks in use.\n",
    "  \"\"\"\n",
    "  del args  # Unused env_config argument passed by ray\n",
    "  return CycleOverBenchmarks(make_env(), train_benchmarks)\n",
    "\n",
    "tune.register_env(\"compiler_gym\", make_training_env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6CVCAIrpyKa4",
    "outputId": "76537a51-91e7-46c5-ab7c-e71122be04c6"
   },
   "outputs": [],
   "source": [
    "# Lets cycle through a few calls to reset() to demonstrate that this environment\n",
    "# selects a new benchmark for each episode.\n",
    "with make_training_env() as env:\n",
    "  env.reset()\n",
    "  print(env.benchmark)\n",
    "  env.reset()\n",
    "  print(env.benchmark)\n",
    "  env.reset()\n",
    "  print(env.benchmark)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IpR63LcOuRRz"
   },
   "source": [
    "## Run the training loop\n",
    "\n",
    "Now that we have the environment set up, let's run a training loop. Here will use RLlib's [Proximal Policy Optimization](https://docs.ray.io/en/master/rllib-algorithms.html#ppo) implementation, and run a very short training loop just for demonstative purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "KYGUxq6GhXZL",
    "outputId": "6ce9e698-6fd2-442d-976d-40c4c0a0902f"
   },
   "outputs": [],
   "source": [
    "import ray\n",
    "from ray.rllib.agents.ppo import PPOTrainer\n",
    "\n",
    "# (Re)Start the ray runtime.\n",
    "if ray.is_initialized():\n",
    "  ray.shutdown()\n",
    "ray.init(include_dashboard=False, ignore_reinit_error=True)\n",
    "\n",
    "tune.register_env(\"compiler_gym\", make_training_env)\n",
    "\n",
    "analysis = tune.run(\n",
    "    PPOTrainer,\n",
    "    checkpoint_at_end=True,\n",
    "    stop={\n",
    "        \"episodes_total\": 500,\n",
    "    },\n",
    "    config={\n",
    "        \"seed\": 0xCC,\n",
    "        \"num_workers\": 1,\n",
    "        # Specify the environment to use, where \"compiler_gym\" is the name we \n",
    "        # passed to tune.register_env().\n",
    "        \"env\": \"compiler_gym\",\n",
    "        # Reduce the size of the batch/trajectory lengths to match our short \n",
    "        # training run.\n",
    "        \"rollout_fragment_length\": 5,\n",
    "        \"train_batch_size\": 5,\n",
    "        \"sgd_minibatch_size\": 5,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CFDxXezNuXdy"
   },
   "source": [
    "## Evaluate the agent\n",
    "\n",
    "After running the training loop we can create a new agent that has exploration disabled, restore it from the training checkpoint, and then use it for running inference tests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4mYoBhSEi26c",
    "outputId": "2dae3ec2-9d0e-4919-86b9-10362d9dd5ee"
   },
   "outputs": [],
   "source": [
    "agent = PPOTrainer(\n",
    "    env=\"compiler_gym\",\n",
    "    config={\n",
    "        \"num_workers\": 1,\n",
    "        \"seed\": 0xCC,\n",
    "        # For inference we disable the stocastic exploration that is used during \n",
    "        # training.\n",
    "        \"explore\": False,\n",
    "    },\n",
    ")\n",
    "\n",
    "# We only made a single checkpoint at the end of training, so restore that. In\n",
    "# practice we may have many checkpoints that we will select from using \n",
    "# performance on the validation set.\n",
    "checkpoint = analysis.get_best_checkpoint(\n",
    "    metric=\"episode_reward_mean\", \n",
    "    mode=\"max\", \n",
    "    trial=analysis.trials[0]\n",
    ")\n",
    "\n",
    "agent.restore(checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "aWeLEVYZjVuM",
    "outputId": "ee293b46-650f-4fbb-8ff4-4414eb4140c5"
   },
   "outputs": [],
   "source": [
    "# Lets define a helper function to make it easy to evaluate the agent's \n",
    "# performance on a set of benchmarks.\n",
    "\n",
    "def run_agent_on_benchmarks(benchmarks):\n",
    "  \"\"\"Run agent on a list of benchmarks and return a list of cumulative rewards.\"\"\"\n",
    "  with make_env() as env:\n",
    "    rewards = []\n",
    "    for i, benchmark in enumerate(benchmarks, start=1):\n",
    "        observation, done = env.reset(benchmark=benchmark), False\n",
    "        while not done:\n",
    "            action = agent.compute_action(observation)\n",
    "            observation, _, done, _ = env.step(action)\n",
    "        rewards.append(env.episode_reward)\n",
    "        print(f\"[{i}/{len(benchmarks)}] {env.state}\")\n",
    "\n",
    "  return rewards\n",
    "\n",
    "# Evaluate agent performance on the validation set.\n",
    "val_rewards = run_agent_on_benchmarks(val_benchmarks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DEc872g10UmF",
    "outputId": "96170474-0742-4e7c-9d53-4d2fc094f9fb"
   },
   "outputs": [],
   "source": [
    "# Evaluate agent performance on the holdout test set.\n",
    "test_rewards = run_agent_on_benchmarks(test_benchmarks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 383
    },
    "id": "Oo3Dn360EE5A",
    "outputId": "55879094-dcc6-4d43-f385-d71db1300efe"
   },
   "outputs": [],
   "source": [
    "# Finally lets plot our results to see how we did!\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def plot_results(x, y, name, ax):\n",
    "  plt.sca(ax)\n",
    "  plt.bar(range(len(y)), y)\n",
    "  plt.ylabel(\"Reward (higher is better)\")\n",
    "  plt.xticks(range(len(x)), x, rotation = 90)\n",
    "  plt.title(f\"Performance on {name} set\")\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2)\n",
    "fig.set_size_inches(13, 3)\n",
    "plot_results(val_benchmarks, val_rewards, \"val\", ax1)\n",
    "plot_results(test_benchmarks, test_rewards, \"test\", ax2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z1lXpJbA04Ji"
   },
   "source": [
    "That's it for this demonstration! Check out the [documentation site](https://facebookresearch.github.io/CompilerGym/) for more details, API reference, and more. If you can encounter any problems, please [file an issue](https://github.com/facebookresearch/CompilerGym/issues)."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "rllib-example.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
